# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import json
import subprocess
import time

import luigi
from luigi.contrib.external_program import ExternalProgramTask
from sqlitedict import SqliteDict

from coin.strategy.mm.simple_sim.runner_slurm import run_simulator_on_slurm


def query_slurm_job(job_id=None):
  fmt = '||'.join(['%%' + c for c in 'AujZTtVMBrRS'])
  query_cmd = ['squeue', '-r', '--format=%s' % fmt]
  if job_id:
    query_cmd += ['-j', str(job_id)]

  try:
    output = subprocess.check_output(query_cmd, stderr=subprocess.PIPE)
    output = output.decode('utf-8')
  except subprocess.CalledProcessError as e:
    stderr = e.stderr.decode('utf-8')
    if (e.returncode == 1 and stderr.strip() == 'slurm_load_jobs error: Invalid job id specified'):
      return None
    raise

  jobs = []
  keys = []
  for idx, line in enumerate(output.strip().split('\n')):
    line = line.strip()
    if idx == 0:
      keys = line.split('||')
      continue
    job = dict(zip(keys, line.split('||')))
    jobs.append(job)

  return jobs or None


def wait_slurm_job(job_id, poll_period=None, deadline=None):
  poll_period = poll_period or datetime.timedelta(seconds=5)

  while True:
    result = query_slurm_job(job_id)
    if result is None:
      return

    assert len(result) == 1
    assert result[0]['JOBID'] == str(job_id), result

    # Slurm job state reference: https://slurm.schedmd.com/squeue.html#lbAG
    state = result[0]['STATE']
    if state == 'COMPLETED':
      return
    elif state in ('BOOT_FAIL',
                   'CANCELLED',
                   'DEADLINE',
                   'FAILED',
                   'NODE_FAIL',
                   'PREEMPTED',
                   'SPECIAL_EXIT',
                   'STOPPED',
                   'TIMEOUT'):
      raise SlurmJobFailedError(result[0])

    if deadline and datetime.datetime.now() > deadline:
      raise TimeoutError()

    time.sleep(poll_period.total_seconds())


class SlurmJobFailedError(RuntimeError):
  def __init__(self, detail):
    super().__init__('Slurm job failed. Slurm Job ID=%s' % detail['JOBID'])
    self.detail = detail


class SqliteTarget(luigi.Target):
  def __init__(self, update_id, db, table=None):
    self.update_id = update_id
    self.db = db
    self.table = table or 'run'

  def _get_table(self, readonly):
    if readonly:
      return SqliteDict(self.db,
                        tablename=self.table,
                        flag='r',
                        encode=json.dumps,
                        decode=json.loads)
    else:
      return SqliteDict(self.db,
                        tablename=self.table,
                        autocommit=True,
                        encode=json.dumps,
                        decode=json.loads)

  def exists(self):
    with self._get_table(readonly=True) as table:
      return self.update_id in table

  def get(self):
    with self._get_table(readonly=True) as table:
      value = table[self.update_id]
      return value[1]

  def touch(self, value=None):
    with self._get_table(readonly=False) as table:
      table[self.update_id] = [str(datetime.datetime.now()), value]

  def reset(self):
    with self._get_table(readonly=False) as table:
      try:
        del table[self.update_id]
      except KeyError:
        pass


class CommandTask(ExternalProgramTask):
  run_id = luigi.Parameter(default=None)
  run_db = luigi.Parameter()
  args = luigi.ListParameter()

  def program_args(self):
    return list(self.args)

  def get_run_id(self):
    return self.run_id if self.run_id is not None else self.task_id

  def run(self):
    super().run()
    self.output().touch()

  def output(self):
    return SqliteTarget(update_id=self.get_run_id(), db=self.run_db)


class _MarketSimLaunchTask(luigi.Task):
  run_id = luigi.Parameter()
  run_db = luigi.Parameter()
  sim_args = luigi.DictParameter()

  def run(self):
    job_id = run_simulator_on_slurm(**self.sim_args)
    self.output().touch(str(job_id))

  def output(self):
    update_id = '%s-launch' % self.run_id
    return SqliteTarget(update_id=update_id, db=self.run_db)


class _MarketSimWaitTask(luigi.Task):
  run_id = luigi.Parameter()
  run_db = luigi.Parameter()
  sim_args = luigi.DictParameter()

  def requires(self):
    return _MarketSimLaunchTask(run_id=self.run_id, run_db=self.run_db, sim_args=self.sim_args)

  def run(self):
    job_id = str(self.input().get())

    try:
      wait_slurm_job(job_id)
    except (SlurmJobFailedError, TimeoutError):
      self.input().reset()
      raise MarketSimRunError(job_id=job_id)

    self.output().touch(job_id)

  def output(self):
    update_id = '%s-done' % self.run_id
    return SqliteTarget(update_id=update_id, db=self.run_db)


class MarketSimTask(luigi.Task):
  run_id = luigi.Parameter()
  run_db = luigi.Parameter()
  sim_args = luigi.DictParameter()

  def requires(self):
    return _MarketSimWaitTask(run_id=self.run_id, run_db=self.run_db, sim_args=self.sim_args)

  def output(self):
    return self.input()


class MarketSimRunError(RuntimeError):
  def __init__(self, job_id):
    super().__init__('MarketSim failed. Slurm Job ID=%s' % str(job_id))
    self.job_id = job_id
